package com.ruanxie.dto.classdetect;

import lombok.Builder;
import lombok.Data;

import java.util.Hashtable;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

//这里detect.py需要约定一些格式
/**
 * 存放下识别结果的关键信息
 *  设计构思：
 *      1. 一开始存下的时候只存输出的字符串
 *      2. 在具体的方法中在对这个简单字符串进行处理
 */
@Data
@Builder
public class YoloResultDTO {
    private List<String> labels;//存放标签信息(这里我修改了下原本的程序)
        //格式："label: 类型下标 class x_center y_center width height"
    private String result;//默认返回的结果//image 3/3 图片路径: 384x640 2 heads, 2 Vpersons,
    // 定义匹配物体类型和数量的正则表达式(初始化一次即可)
    static Pattern detectionPattern = Pattern.compile("(\\d+)\\s([a-zA-Z]+), ");

    /**
     * result字符串对应的格式（yolov5默认）
     * 这里用正则处理下
     * @return
     */
    public Map<String,Object> getDetectResult(){
        Map<String,Object> detectResult=new Hashtable<>();
        // 查找检测的类型和数量
        Matcher matcher = detectionPattern.matcher(result);
        while (matcher.find()) {
            // group(1) 是数量，group(2) 是类别
            Integer count =Integer.valueOf(matcher.group(1));
            String type = matcher.group(2);
//                System.out.println("检测到 " + count + " 个 " + type);
            if ("Vpersons".equals(type)){
                type="Vperson";
            }
            detectResult.put(type,count);
        }
        return detectResult;
    }
}
